# -*- coding: utf-8 -*-
from bs4 import BeautifulSoup
import requests
import re

class CWE_INHERITANCE:
    def __init__(self,cweid=None,group=None) -> None:
        self.child_list=[]
        self.cweid=cweid
        self.group=group
        
    def add_child(self,child):
        self.child_list.append(child)

    def get_child_list(self):
        return self.child_list

    def get_cweid(self):
        return self.cweid

    def auto_set_child_list(self):
        id=self.group.attrs['id']
        group_list=self.group.select(f"div[id='{id}'] > div[class='collapseblock'] > div[class='group']")
        if len(group_list)==0:
            self.child_list=[]
            return
        for group in group_list:
            cweid=self.get_cweid(group)
            self.add_child(CWE_INHERITANCE(cweid,group))
        for child in self.child_list:
            child.auto_set_child_list()

    def get_cweid(self,group):
        '''获取group下第一个节点的cwe编号'''
        id=group.attrs['id']
        cwe_info=group.select(f"div[id='{id}'] > span[class='graph_title'] > span > a > span ")[0].text
        cweid='CWE-'+re.search(r' - \((\d+)\)',cwe_info,flags=0).groups()[0]
        return cweid

    def get_childs_cweid(self):
        childs_cweid=[]
        for child in self.child_list:
            cweid=child.cweid
            childs_cweid.append(cweid)
            childs_cweid+=child.get_childs_cweid()
        return childs_cweid
    
    def get_childs(self):
        childs=[]
        for child in self.child_list:
            childs.append(child)
            childs+=child.get_childs()
        return childs

    def get_relations(self,cweida,cweidb):
        '''分析cweida是cweidb的什么关系,返回结果为ancestor说明cweida是cweidb的祖辈,有四种结果:ancestor descendant equal unrelated 分别表示祖辈、子辈、相同、无辈分关系,为图方便此方法仅root节点可正常使用'''
        if cweida==cweidb:
            return 'equal'
        root_childs=self.get_childs()
        root_childs_cweid=self.get_childs_cweid()
        if cweida not in root_childs_cweid:
            print(f'{cweida}在cwe1000中找不到')
            return 'unrelated'
        if cweidb not in root_childs_cweid:
            print(f'{cweidb}在cwe1000中找不到')
            return 'unrelated'
        index=0
        # 存在一个CWE变化在CWE1000中出现多次的情况，因此需要循环一下
        for i in range(root_childs_cweid.count(cweida)):
            cweida_index=root_childs_cweid.index(cweida,index)
            cweida_child=root_childs[cweida_index]
            cweida_child_childs_cweid=cweida_child.get_childs_cweid()
            if cweidb in cweida_child_childs_cweid:
                return 'ancestor'
            index=cweida_index+1
        index=0
        for i in range(root_childs_cweid.count(cweidb)):
            cweidb_index=root_childs_cweid.index(cweidb,index)
            cweidb_child=root_childs[cweidb_index]
            cweidb_child_childs_cweid=cweidb_child.get_childs_cweid()
            if cweida in cweidb_child_childs_cweid:
                return 'descendant'
            index=cweidb_index+1
        return 'unrelated'

def get_cweid(group):
    '''获取group下第一个节点的cwe编号'''
    id=group.attrs['id']
    cwe_info=group.select(f"div[id='{id}'] > span[class='graph_title'] > span[class='Primary'] > a > span[class='cweid Primary']")[0].text
    cweid='CWE-'+re.search(r' - \((\d+)\)',cwe_info,flags=0).groups()[0]
    return cweid

def get_cwe_root():
    html=requests.get(url='https://cwe.mitre.org/data/definitions/1000.html').content
    soup = BeautifulSoup(html, 'html.parser')
    group_list=soup.select("div[class='tabledetail'] > div[class='indent'] > div > div > div[class='group']")
    cwe_root=CWE_INHERITANCE()
    for group in group_list:
        cweid=get_cweid(group)
        cwe_root.add_child(CWE_INHERITANCE(cweid,group))
    for child in cwe_root.get_child_list():
        child.auto_set_child_list()
    return cwe_root

# while(1):
#     cweida=input()
#     cweidb=input()
#     print(cwe_root.get_relations(cweida,cweidb))